import networkx as nx
import matplotlib
matplotlib.use(backend="TkAgg")
import matplotlib.pyplot as plt
import numpy as np

# 使用之前定义的邻接矩阵
adj = np.array([
    [0,1,1,0],
    [1,0,1,1],
    [1,1,0,1],
    [0,1,1,0]
])

# 创建无向图
G = nx.Graph()

# 添加节点
num_nodes = len(adj)
G.add_nodes_from(range(num_nodes))

# 添加边
for i in range(num_nodes):
    for j in range(i+1, num_nodes):  # 只遍历上三角，避免重复
        if adj[i, j] == 1:
            G.add_edge(i, j)

# 计算度数和平稳分布
degrees = adj.sum(axis=1)
s_rw = degrees / degrees.sum()

# 设置图形布局
plt.figure(figsize=(10, 8))

# 使用spring布局让图形更美观
pos = nx.spring_layout(G, seed=42)

# 绘制图形
# 节点大小与度数成正比
node_sizes = [s * 3000 for s in s_rw]  # 缩放因子使大小合适

# 节点颜色与平稳分布概率成正比
node_colors = s_rw

# 绘制节点
nodes = nx.draw_networkx_nodes(G, pos,
                              node_size=node_sizes,
                              node_color=node_colors,
                              cmap=plt.cm.Blues,
                              alpha=0.8)

# 绘制边
nx.draw_networkx_edges(G, pos, alpha=0.5, width=2)

# 添加节点标签
labels = {i: f'Node {i}\nDeg: {degrees[i]}\nπ: {s_rw[i]:.2f}' for i in range(num_nodes)}
nx.draw_networkx_labels(G, pos, labels, font_size=10)

# 添加边标签（转移概率）
edge_labels = {}
for i, j in G.edges():
    # 两个方向的转移概率
    p_ij = 1/degrees[i]  # i->j的概率
    p_ji = 1/degrees[j]  # j->i的概率
    edge_labels[(i, j)] = f'{p_ij:.2f}/{p_ji:.2f}'

nx.draw_networkx_edge_labels(G, pos, edge_labels, font_size=8)

plt.title("Random Walk on Undirected Graph\n(Node size/color = Stationary probability)")
plt.axis('off')  # 关闭坐标轴

# 添加颜色条
plt.colorbar(nodes, label='Stationary Probability')

plt.tight_layout()
plt.show()

# 额外：绘制转移概率矩阵的热力图
plt.figure(figsize=(8, 6))
P_rw = (adj.T / degrees).T  # 计算转移矩阵
plt.imshow(P_rw, cmap='Blues', interpolation='nearest')
plt.colorbar(label='Transition Probability')
plt.title("Transition Probability Matrix P")
plt.xticks(range(num_nodes))
plt.yticks(range(num_nodes))
for i in range(num_nodes):
    for j in range(num_nodes):
        plt.text(j, i, f'{P_rw[i, j]:.2f}',
                ha='center', va='center',
                color='white' if P_rw[i, j] > 0.5 else 'black')
plt.xlabel('To State')
plt.ylabel('From State')
plt.tight_layout()
plt.show()

# 打印详细统计信息
print("Graph Details:")
print(f"Nodes: {list(G.nodes())}")
print(f"Edges: {list(G.edges())}")
print(f"Degrees: {degrees}")
print(f"Stationary distribution: {s_rw}")
print("\nTransition Matrix P:")
print(P_rw)